
# 一维谐振子求基态能量
import numpy as np
import matplotlib.pyplot as plt

# trial wave: psi = e^{ -alpha x^2 }
# H = - d^2/dx^2 + x*x
import scipy.integrate

def psi(x,alpha):
    return x*np.exp( -alpha * x * x )

#help(scipy.integrate.quad); exit(1)
def plotpsi2(alpha):
    C = scipy.integrate.quad(psi,-np.inf,np.inf,args=(2*alpha,))[0]
    x = np.arange(-3,3,0.1)
    y = [  psi(t,alpha) * psi(t,alpha) / C for t in x ]
    plt.plot(x,y)
#plotpsi2(0.4); exit(1)

# d/dx psi = -2 alpha x e^{- alpha * x * x }
# d^2 /d x^2 psi = -2 alpha + 4 * alpha^2 x^2 e^{-alpha x^2}
# ( - d^2/dx^2 + x*x ) psi / psi =  2 alpha + (- 4 alpha^2 + 1) x^2
def EL( alpha, x):
    #return 2*alpha + (1-4*alpha*alpha)*x*x
    return 6 * alpha + (1 - 4 * alpha * alpha) * x * x

# if alpha = 0.5, EL(alpha, x) = 1
# if alpha = 0.6, EL(alpha, x) = 1.2 - 0.44 * x*x

def Accept( x1, x2, alpha, psi ):
    if( abs(x2)>3 ): return 0
    p = ( psi(x2,alpha) / psi(x1,alpha) )**2
    #print("p=",p)
    if p>1: return 1
    else: return p

def tryonestep(x,nwalker,h,alpha):
    for j in range(nwalker):
        step = h * ( 2*np.random.randint(0,2) - 1 )
        A = Accept(x[j], x[j] + step, alpha, psi)
        if (np.random.random() < A):
            x[j] += step

def ELsampling( alpha, nwalker, psi, EL, a, b, h, nstep ):
    x = (b-a)*np.random.random(nwalker) + a # 初始位置
    for i in range(nstep):
        tryonestep(x, nwalker, h, alpha)
    #plotpsi2(alpha);
    #plt.hist(x, bins=np.arange(-3.1,3.1,0.2), rwidth=0.8, density="True" ); plt.show()
    #print( "alpha = ",alpha, " <E>_L = ", np.average( [ EL(alpha,t) for t in x ] ) )
    return [ EL(alpha,t) for t in x ]

aveE = []; aveE2 = []
Alpha = np.arange(0.4,0.6,0.02)
for alpha in Alpha:
    ELsample = ELsampling(alpha, 1000, psi, EL, a=-3, b=3, h=0.2, nstep=1000 )
    aveE.append( np.average(ELsample) )
    aveE2.append( np.average([t*t for t in ELsample ]) )
sigma2 = aveE2 - np.array([t*t for t in aveE])
plt.xlabel(r"$\alpha$",fontsize=15)
plt.ylabel(r"$\sigma^2(\alpha)$",fontsize=15)
plt.plot(Alpha, sigma2)
plt.savefig("VMC一维谐振子.png")
plt.show()